#!/usr/bin/python

# sf-muni.py: Provides a custom loader for the SF Muni. 

# Copyright 2011 Portland Transport

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

#   http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# Authors:
# Matt Conway: main code

# The GTFS file URL is:
# http://www.sfmta.com/transitdata/google_transit.zip
# Note you must pass this in using the -g option, as well as passing in
# -a sf-muni
# Note there is a license agreement to use the above URL.

# Don't do tags as IDs for SF Muni, it expects to find the stop_ids in
# the stop_code attribute of the NextBus stops

# Basically, the SF Muni's GTFS file uses stop ids which are the same as
# nextbuses stop_codes

import transitfeed
import nextbusloader

class MuniLoader(nextbusloader.NextBusLoader):
    def loadAllStops(self):
        nbstops_unindexed = self._getStopsFromNextBus()

       # First, build an index by stop_code of the NextBus stops
        nbstops = dict([(stop.stop_code, stop) for stop in nbstops_unindexed])

        # This one will set self.agency more accurately
        gtfsstops = self._getStopsFromGTFS(
            self.options.gtfsUrl or 'http://www.sfmta.com/transitdata/google_transit.zip',
            url=not self.options.localgtfs) # not turns False into True, vice-versa
        
        self.agency.avl_service = 'NextBus'

        # Now, correlate them.
        # Use this to add the NextBus ID to GTFS
        def addNextBusID(stop):
            # The nextbus ID is indexed by the NB stop_tag, which is the same
            # as GTFS stop_id.
            try:
                stop.avl_id = nbstops[stop.stop_id].stop_id
            except KeyError:
                print 'Could not find stop %s (%s) in NextBus feed - perhaps it is orphaned? Ignoring.' % (stop.stop_id, self.agency.agency_name + stop.stop_name)
                return None # We will remove all of the Nones in a moment

            return stop

        # Remove all of the None stops
        self.stops = [v for v in map(addNextBusID, gtfsstops) if v != None]

if __name__ == '__main__':
    nextbusloader.main(loaderClass=MuniLoader)

